# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/3 上午11:26
@Author yzpang
@Function: 
"""
from typing import Any

import pydantic
from pydantic import BaseModel
from modelserver.configs.server_config import MULTI_CLASS
from modelserver.configs.model_config import DEFAULT_DEVICE


class BaseRequest(BaseModel):
    """基础请求"""
    description: str = pydantic.Field(None, description='Request description')
    body: Any = pydantic.Field(None, description='Request body')

    class Config:
        json_schema_extra = {
            "example": {
                "description": "request to load model",
                "body": {
                },
            }
        }


class ModelMessage(BaseModel):
    """模型管理请求对象"""
    modelName: str = pydantic.Field(..., description='Model name')
    modelPath: str = pydantic.Field(None, description='Model path')
    modelFrame: str = pydantic.Field("TransformersModelWorker", description='Model frame')
    device: str = pydantic.Field(DEFAULT_DEVICE, description='Device')
    classificationType: str = pydantic.Field(MULTI_CLASS, description='Classification type')

    class Config:
        json_schema_extra = {
            "example": {
                "modelName": "bert-base-cased",
                "modelPath": "bert-base-cased",
                "classificationType": MULTI_CLASS,
            }
        }


class InferenceMessage(BaseRequest):
    """推理请求对象"""
    modelName: str = pydantic.Field(..., description='Model name')
    text: str = pydantic.Field(..., description='Inference message')

    class Config:
        json_schema_extra = {
            "example": {
                "modelName": "bert-base-cased",
                "text": "汽车尾气排放不畅",
            }
        }


class EvaluateMessage(BaseRequest):
    """模型评估请求对象"""
    modelPath: str = pydantic.Field(None, description='Model path')
    labelPath: str = pydantic.Field(..., description='Label path')
    dataPath: str = pydantic.Field(None, description='Data path')
    classificationType: str = pydantic.Field(MULTI_CLASS, description='Classification type')

    class Config:
        json_schema_extra = {
            "example": {
                "modelPath": "BERT/bert-base-cased",
                "labelPath": "label.txt",
                "dataPath": "data.txt",
                "classificationType": MULTI_CLASS,
            }
        }


class TrainMessage(BaseRequest):
    """模型训练请求对象"""
    modelName: str = pydantic.Field(..., description='Model name')  # 必填
    modelPath: str = pydantic.Field(None, description='Model path')
    labelPath: str = pydantic.Field(None, description='Label path')
    dataPath: str = pydantic.Field(None, description='Data path')
    evalPath: str = pydantic.Field(None, description='Eval Path')
    outputPath: str = pydantic.Field(None, description='Output Path')
    classificationType: str = pydantic.Field(MULTI_CLASS, description='Classification type')
    extra: dict = pydantic.Field(None, description="Extra args")

    class Config:
        json_schema_extra = {
            "example": {
                "modelName": "bert-base-cased",
                "modelPath": "BERT/bert-base-cased",
                "labelPath": "label.txt",
                "dataPath": "data.txt",
                "evalPath": "eval.txt",
                "outputPath": "output",
                "classificationType": MULTI_CLASS,
                "extra": {}
            }
        }


